import pymysql
import pandas as pd

class MysqlUtils(object):
    """MySQL数据库工具类"""
    
    def __init__(self):
        """初始化数据库连接"""
        self.conn = pymysql.connect(
            host='127.0.0.1',
            user='root',
            passwd='sjk1234',
            db='tushare',
            port=3306,
            charset='utf8'
        )

class Classification(object):
    """财务数据分析分类工具类"""
    
    def __init__(self):
        pass
    
    def get_fina_indicator(self, conn):
        """
        获取财务指标数据
        参数:
            conn: 数据库连接对象
        返回:
            DataFrame格式的财务指标数据
        """
        cursor = conn.cursor(cursor=pymysql.cursors.DictCursor)
        sql = """
            SELECT ts_code, ann_date, eps, total_revenue_ps, undist_profit_ps, 
                   gross_margin, fcff, fcfe, tangible_asset, bps, 
                   grossprofit_margin, npta
            FROM fina_indicator
        """
        cursor.execute(sql)
        ret = cursor.fetchall()
        
        df = pd.DataFrame(ret)
        # 数据清洗 - 删除关键字段为空的行
        df = df.dropna(subset=['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin'])
        # 重置索引
        df = df.reset_index(drop=True)
        print(df.head())
        
        return df
    
    def get_daily(self, conn, df):
        """
        获取每日交易数据并与财务指标合并
        参数:
            conn: 数据库连接对象
            df: 包含财务指标的DataFrame
        返回:
            合并后的DataFrame并保存为CSV文件
        """
        cursor = conn.cursor(cursor=pymysql.cursors.DictCursor)
        new_list = []
        
        for i in range(len(df['ts_code'])):
            try:
                # 获取每只股票的交易数据
                sql = f"SELECT trade_date, close FROM daily_data WHERE ts_code = '{df['ts_code'][i]}'"
                cursor.execute(sql)
                ret = cursor.fetchall()
                df1 = pd.DataFrame(ret)
                
                if len(df1) > 0:
                    # 计算关键指标
                    max_close = df1['close'].max()
                    min_close = df1['close'].min()
                    the_close = df1['close'].iloc[-1]  # 获取最新收盘价
                    
                    # 组合财务指标和交易数据
                    new_list.append({
                        'ts_code': df['ts_code'][i],
                        'ann_date': df['ann_date'][i],
                        'max_close': max_close,
                        'min_close': min_close,
                        'the_close': the_close,
                        'eps': df['eps'][i],
                        'total_revenue_ps': df['total_revenue_ps'][i],
                        'undist_profit_ps': df['undist_profit_ps'][i],
                        'gross_margin': df['gross_margin'][i],
                        'fcff': df['fcff'][i],
                        'fcfe': df['fcfe'][i],
                        'tangible_asset': df['tangible_asset'][i],
                        'bps': df['bps'][i],
                        'grossprofit_margin': df['grossprofit_margin'][i],
                        'npta': df['npta'][i]
                    })
                    
            except Exception as e:
                print(f"处理股票{df['ts_code'][i]}时出错: {e}")
        
        # 转换为DataFrame并保存
        df2 = pd.DataFrame(new_list)
        df2.to_csv('daily.csv', index=False)
        print("数据已保存到daily.csv")
        return df2

if __name__ == '__main__':
    # 主程序入口
    mu = MysqlUtils()  # 初始化数据库连接
    cl = Classification()  # 初始化分析工具
    
    # 获取财务指标数据
    fina_df = cl.get_fina_indicator(mu.conn)
    
    # 获取并合并每日交易数据
    if not fina_df.empty:
        cl.get_daily(mu.conn, fina_df)